import sys

# 直接复用Pipeline实例
sys.path.append("./")
pipeline = __import__("03-token-classification")

import gradio as gr


# 定义分类函数
def classify_tokens(text):
    # 使用Pipeline进行分类
    result = pipeline.nlp(text)

    # 提取分类标签和词元
    output = []
    for entity in result:
        word = entity['word']
        label = entity['entity']
        output.append((word, label))

    # 返回格式化后的结果
    return output


# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# 词元分类器")
    gr.Markdown(
        "这是一个基于Transformers框架的命名实体识别工具。您可以输入任意文本，点击“提交”按钮后，系统将自动标记文本中的实体。")

    with gr.Row():
        input_text = gr.Textbox(placeholder="请输入要分类的文本...", label="输入文本")

    with gr.Row():
        submit_button = gr.Button("提交")

    with gr.Row():
        output_label = gr.HighlightedText(label="分类结果")

    # 设置按钮点击事件，触发分类函数
    submit_button.click(classify_tokens, inputs=input_text, outputs=output_label)

# 启动Gradio应用
if __name__ == "__main__":
    demo.launch()
